//===--- StackNesting.h - Utility for stack nesting -------------*- C++ -*-===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

#ifndef POLARPHP_PIL_OPTIMIZER_UTILS_STACKNESTING_H
#define POLARPHP_PIL_OPTIMIZER_UTILS_STACKNESTING_H

#include "polarphp/pil/lang/PILInstruction.h"
#include "llvm/ADT/SmallBitVector.h"

#include <vector>

namespace polar {

/// A utility to correct the nesting of stack allocating/deallocating
/// instructions.
///
/// This utility is useful for optimizations which create stack allocations
/// without caring about correct nesting with existing allocations. This may
/// result in code like
/// \code
///   %1 = alloc_stack
///     ...
///   %2 = alloc_stack
///     ...
///   dealloc_stack %1
///     ...
///   dealloc_stack %2
/// \endcode
///
/// The StackNesting utility is able to correct the code:
/// \code
///   %1 = alloc_stack
///     ...
///   %2 = alloc_stack
///     ...
///   dealloc_stack %2
///   dealloc_stack %1
/// \endcode
///
class StackNesting {

   typedef SmallBitVector BitVector;

   /// Data stored for each block (actually for each block which is not dead).
   struct BlockInfo {

      /// Back-link to the block.
      PILBasicBlock *Block;

      /// The cached list of successors.
      llvm::SmallVector<BlockInfo *, 8> Successors;

      /// The list of stack allocating/deallocating instructions in the block.
      llvm::SmallVector<PILInstruction *, 8> StackInsts;

      /// The bit-set of alive stack locations at the block entry.
      BitVector AliveStackLocsAtEntry;

      /// The bit-set of alive stack locations at the block exit.
      BitVector AliveStackLocsAtExit;

      BlockInfo(PILBasicBlock *Block) : Block(Block) { }
   };

   /// Data stored for each stack location (= allocation).
   ///
   /// Each stack location is allocated by a single allocation instruction.
   struct StackLoc {
      StackLoc(SingleValueInstruction *Alloc) : Alloc(Alloc) { }

      /// Back-link to the allocation instruction.
      SingleValueInstruction *Alloc;

      /// Bit-set which represents all alive locations at this allocation.
      /// It obviously includes this location itself. And it includes all "outer"
      /// locations which surround this location.
      BitVector AliveLocs;
   };

   /// Mapping from stack allocations (= locations) to bit numbers.
   llvm::DenseMap<SingleValueInstruction *, unsigned> StackLoc2BitNumbers;

   /// The list of stack locations. The index into this array is also the bit
   /// number in the bit-sets.
   llvm::SmallVector<StackLoc, 8> StackLocs;

   /// Block data for all (non-dead) blocks.
   std::vector<BlockInfo> BlockInfos;

public:

   /// The possible return values of correctStackNesting().
   enum class Changes {
      /// No changes are made.
         None,

      /// Only instructions were inserted or deleted.
         Instructions,

      /// Instructions were inserted or deleted and new blocks were inserted.
         CFG
   };

   StackNesting() { }

   /// Performs correction of stack nesting by moving stack-deallocation
   /// instructions down the control flow.
   ///
   /// Returns the status of what changes were made.
   Changes correctStackNesting(PILFunction *F);

   /// For debug dumping.
   void dump() const;

   static void dumpBits(const BitVector &Bits);

private:
   /// Initializes the data structures.
   void setup(PILFunction *F);

   /// Solves the dataflow problem.
   ///
   /// Returns true if there is a nesting of locations in any way, which can
   /// potentially in the wrong order.
   bool solve();

   /// Insert deallocation instructions for all locations which are alive before
   /// the InsertionPoint (AliveBefore) but not alive after the InsertionPoint
   /// (AliveAfter).
   ///
   /// Returns true if any deallocations were inserted.
   bool insertDeallocs(const BitVector &AliveBefore, const BitVector &AliveAfter,
                       PILInstruction *InsertionPoint,
                       Optional<PILLocation> Location);

   /// Reeturns the location bit number for a stack allocation instruction.
   int bitNumberForAlloc(PILInstruction *AllocInst) {
      assert(AllocInst->isAllocatingStack());
      return StackLoc2BitNumbers[cast<SingleValueInstruction>(AllocInst)];
   }

   /// Reeturns the location bit number for a stack deallocation instruction.
   int bitNumberForDealloc(PILInstruction *DeallocInst) {
      assert(DeallocInst->isDeallocatingStack());
      auto *AllocInst = cast<SingleValueInstruction>(DeallocInst->getOperand(0));
      return bitNumberForAlloc(AllocInst);
   }

   /// Insert deallocations at block boundaries.
   Changes insertDeallocsAtBlockBoundaries();

   /// Modifies the PIL to end up with a correct stack nesting.
   ///
   /// Returns the status of what changes were made.
   bool adaptDeallocs();
};

} // end namespace polar

#endif // POLARPHP_PIL_OPTIMIZER_UTILS_STACKNESTING_H
